"""
对 flask-sqlalchemy 进行二次开发, 实现读写分离



基本实现思路:
    实现自定义的 session类, 继承 SignallingSession类
        重写 get_bind方法, 根据读写需求选择对应的数据库地址
    实现自定义的 SQLAlchemy类, 继承 SQLAlchemy类
        重写 create_session方法, 在内部使用自定义的 Session类
"""
import random
from flask import Flask
from flask_sqlalchemy import SQLAlchemy, SignallingSession, get_state
from sqlalchemy import orm
from sqlalchemy.sql.dml import UpdateBase
from sqlalchemy.orm import Session as SessionBase

app = Flask(__name__)

# 设置单个数据库URI (用于建表)
app.config['SQLALCHEMY_DATABASE_URI'] = 'mysql://root:mysql@192.168.19.137:3306/test31'

# 设置多个数据库的URI (用于数据操作)
app.config['SQLALCHEMY_BINDS'] = {
    'master': 'mysql://root:mysql@192.168.19.137:3306/test31',
    'slave1': 'mysql://root:mysql@192.168.19.137:8306/test31',
    'slave2': 'mysql://root:mysql@192.168.19.137:3306/test31'
}

# 其他配置
app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = False
app.config['SQLALCHEMY_ECHO'] = True


# 1. 自定义Session类, 继承SignallingSession, 并重写get_bind方法
class RoutingSession(SignallingSession):
    def __init__(self, *args, **kwargs):
        super(RoutingSession, self).__init__(*args, **kwargs)
        # 每个Session(请求), 随机一次从库
        self.slave = random.choice(['slave1', 'slave2'])

    def get_bind(self, mapper=None, clause=None):
        """每次数据库操作(增删改查及事务操作)都会调用该方法, 来获取对应的数据库引擎(访问的数据库)"""

        state = get_state(self.app)

        if mapper is not None:  # 如果该操作中涉及的模型类和数据表建立了映射
            try:
                # SA >= 1.3
                persist_selectable = mapper.persist_selectable
            except AttributeError:
                # SA < 1.3
                persist_selectable = mapper.mapped_table

            info = getattr(persist_selectable, 'info', {})
            bind_key = info.get('bind_key')  # 查询模型类是否指定了访问的数据库

            if bind_key is not None:  # 如果该模型类已指定数据库, 使用指定的数据库
                return state.db.get_engine(self.app, bind=bind_key)

        # 如果模型类没有指定要访问的库, 则访问默认的数据库(SQLALCHEMY_DATABASE_URI指定的库)
        # return SessionBase.get_bind(self, mapper, clause)

        # 如果模型类没有指定要访问的库, 则根据读写行为访问对应的主从数据库
        # isinstance(对象, 类)  判断该对象是否为指定的类型  True/False
        if self._flushing or isinstance(clause, UpdateBase):  # 如果模型类未指定数据库, 判断是否为写操作(增删改操作)
        # if self._flushing:
            print('写操作')
            return state.db.get_engine(self.app, bind='master')
        else:

            print('读操作: ', self.slave)
            return state.db.get_engine(self.app, bind=self.slave)


# 2. 自定义SQLALchemy类, 重写create_session方法,在内部使用自定义的 Session类
class RoutingSQLAlchemy(SQLAlchemy):
    def create_session(self, options):
        return orm.sessionmaker(class_=RoutingSession, db=self, **options)


# 创建组件对象
db = RoutingSQLAlchemy(app)
# print(db.session)   # 这样使用的session就变为RoutingSession


# 构建模型类
class User(db.Model):
    __tablename__ = 't_user'
    # __bind_key__ = 'master'  # 指定模型访问的数据库
    id = db.Column(db.Integer, primary_key=True)
    name = db.Column('username', db.String(20), unique=True)
    age = db.Column(db.Integer, default=0, index=True)


@app.route('/')
def index():
    """增加数据"""

    print('---读-----------')

    users = User.query.all()
    for user in users:
        print(user.id, user.name, user.age)

    print('---写-----------')

    user1 = User(name='zs', age=20)
    db.session.add(user1)
    db.session.commit()

    print('---读-----------')

    users = User.query.all()
    for user in users:
        print(user.id, user.name, user.age)

    return "index"


if __name__ == '__main__':
    # 重置所有继承自db.Model的表
    # 如果模型类没有设置__bind_ky__属性(指定对应的数据库), 则DDL操作 根据SQLALCHEMY_DATABASE_URI 指定的数据库进行处理
    db.drop_all()
    db.create_all()
    app.run(debug=True, host='0.0.0.0')
